
###############################
## TODO: 目标不使用任何库，用C语言实现如下python函数功能

import torch
import numpy

class Detection:

    def __init__(self):
        #self.__anchors = torch.Tensor( [(1.25, 1.625), (2.0, 3.75), (4.125, 2.875)])             # Anchors for small obj
        #self.__anchors = torch.Tensor([(1.875, 3.8125), (3.875, 2.8125), (3.6875, 7.4375)]),    # Anchors for medium obj
        self.__anchors = torch.Tensor([(3.625, 2.8125), (4.875, 6.1875), (11.65625, 10.1875)])   # Anchors for big obj

        #self.__stride = 8         # Anchors for small obj
        #self.__stride = 16        # Anchors for medium obj
        self.__stride = 32     # Anchors for big obj


        self.__nA = len(self.__anchors) # = 3
        self.__nC = 20           # class num for voc dataset
        

        self.training = True

    ## TODO: 传入参数p的数据类型， 维度, 参考？ 能否提供验证的输入输出集?
    def forward(self, p):
        bs, nG = p.shape[0], p.shape[-1]
        p = p.view(bs, self.__nA, 5 + self.__nC, nG, nG).permute(0, 3, 4, 1, 2)

        p_de = self.__decode(p.clone())

        return (p, p_de)


    def __decode(self, p):
        batch_size, output_size = p.shape[:2]

        device = p.device
        stride = self.__stride
        anchors = (1.0 * self.__anchors).to(device)

        conv_raw_dxdy = p[ :, :, :, :, 0:2]
        conv_raw_dwdh = p[ :, :, :, :, 2:4]
        conv_raw_conf = p[ :, :, :, :, 4:5]
        conv_raw_prob = p[ :, :, :, :, 5:]

        y = torch.arange(0, output_size).unsqueeze(1).repeat(1, output_size)
        x = torch.arange(0, output_size).unsqueeze(0).repeat(output_size, 1)
        grid_xy = torch.stack([x, y], dim = -1)
        grid_xy = grid_xy.unsqueeze(0).unsqueeze(3).repeat(batch_size, 1, 1, 3, 1).double().to(device)

        pred_xy = (torch.sigmoid(conv_raw_dxdy) + grid_xy) * stride
        pred_wh = (torch.exp(conv_raw_dwdh) * anchors) * stride
        pred_xywh = torch.cat([pred_xy, pred_wh], dim = -1)
        pred_conf = torch.sigmoid(conv_raw_conf)
        pred_prob = torch.sigmoid(conv_raw_prob)
        pred_bbox = torch.cat([pred_xywh, pred_conf, pred_prob], dim = -1)

        ## TODO: self.training 是什么？
        return pred_bbox.view(-1, 5 + self.__nC) if not self.training else pred_bbox


    def load_p(self, file_name):
        # p的shape small [1, 75, 64, 64]  medium[1, 75, 32, 32]  big[1, 75, 16, 16]

        with open(file_name, 'r') as f1:
            list1 = f1.readlines()
        for i in range(0, len(list1)):
            list1[i] = float(list1[i].rstrip('\n'))
        #print(list1)
        p = torch.Tensor(list1)
        p = p.view(1, 75, 16, 16)
        
        return p



def main():
   dect =  Detection()
   p = dect.load_p("/home/ningjian/Code/torch_forward_c/data/big_in.txt")
   result = dect.forward(p)
   print(result[1])




if __name__ == '__main__':
    main()
